{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "96b2b1e4",
            "metadata": {},
            "source": [
                "# Async TreeSummarizeQuery Demo"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "9331cfeb",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# NOTE: This is ONLY necessary in jupyter notebook.\n",
                "# Details: Jupyter runs an event-loop behind the scenes. \n",
                "#          This results in nested event-loops when we start an event-loop to make async queries.\n",
                "#          This is normally not allowed, we use nest_asyncio to allow it for convenience.  \n",
                "import nest_asyncio\n",
                "nest_asyncio.apply()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "8a1d2821",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "import time\n",
                "from gpt_index import GPTListIndex, SimpleDirectoryReader"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "4466dec2",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "query_str = \"What is Paul Graham's biggest achievement?\""
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "6948df36",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# load documents\n",
                "documents = SimpleDirectoryReader('../paul_graham_essay/data').load_data()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "466c3892",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "index = GPTListIndex.from_documents(documents)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "2d9115d1",
            "metadata": {},
            "source": [
                "#### By default, generate a response through hierarchical tree summarization (i.e., `response_mode=tree_summarize`) makes blocking LLM calls"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d9ef0fef",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "start_time = time.perf_counter()\n",
                "index.query(query_str, response_mode='tree_summarize')\n",
                "elapsed_time = time.perf_counter() - start_time"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "9392d573",
            "metadata": {},
            "source": [
                "It takes ~21s to generate a response through hierarchical tree summarization (i.e., `response_mode=tree_summarize`)."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "43fa6067-493e-42e5-8f6d-2c65a7a67e75",
            "metadata": {},
            "source": [
                "#### Option 1: Running `aquery` (async query call) will take advantage of async LLM calls"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f29da0b6-7454-4661-88ab-9ebc80533743",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "import asyncio\n",
                "start_time = time.perf_counter()\n",
                "task = index.aquery(query_str, response_mode='tree_summarize')\n",
                "asyncio.run(task)\n",
                "elapsed_time = time.perf_counter() - start_time"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e9b0f10f-8d74-4db4-a962-b2d00c575762",
            "metadata": {
                "tags": []
            },
            "source": [
                "It takes ~6.8s to generate a response through `aquery` (i.e., `response_mode=tree_summarize`)."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "474f82d1",
            "metadata": {},
            "source": [
                "#### Option 2: Pass in `use_async=True` to enable asynchronous LLM calls within a synchronous `query`\n",
                "\n",
                "This approach makes a synchronous `query` calls, but runs async tasks during the \"tree_summarize\" operation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "78a02987",
            "metadata": {
                "scrolled": true,
                "tags": []
            },
            "outputs": [],
            "source": [
                "start_time = time.perf_counter()\n",
                "index.query(query_str, response_mode='tree_summarize', use_async=True)\n",
                "elapsed_time = time.perf_counter() - start_time"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "23469128",
            "metadata": {},
            "source": [
                "It takes ~6.9s to generate a response through hierarchical tree summarization (i.e., `response_mode=tree_summarize`)."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "llama_index",
            "language": "python",
            "name": "llama_index"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.10"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}